# -*- coding: utf-8 -*-
# time: 2025/4/12 08:49
# file: 儿科问诊数据集.py
# author: hanson

# 加载远程医疗数据集‌:ml-citation{ref="1" data="citationList"}
from datasets import load_dataset
from transformers import AutoTokenizer

dataset = load_dataset("ticoAg/Chinese-medical-dialogue", split="train")
# 转换为PyTorch Dataset格式（可选）‌:ml-citation{ref="4" data="citationList"}
# 查看前3条样本‌:ml-citation{ref="1,8" data="citationList"}
print(dataset)  # 输出字段结构及示例数据
print()
ds = dataset.select(range(10))
#测试数据
for i in ds:
    print(i)




# 2. 数据预处理
def preprocess_data( dataset):
    def format_example(ex):
        # 清理数据中的None值
        instruction = ex["instruction"] or ""
        input_text = ex["input"] or ""
        output = ex["output"] or ""

        text = f"Instruction: {instruction.strip()}\n"
        if input_text:
            text += f"Input: {input_text.strip()}\n"
        text += f"Output: {output.strip()}"
        return {"text": text}

    formatted = dataset.map(format_example, remove_columns=dataset.column_names)
    return formatted,dataset

if __name__ == '__main__':
    formatted,dataset2 = preprocess_data(ds)
    for i in formatted:
        print(i)

    # for i in dataset2:
    #     print(i)
    #     print(i["instruction"] +"_"+i["input"]+"_"+i["output"])
